import numpy as np
import cv2
from scipy.optimize import least_squares

def euler_to_matrix(euler):
    """欧拉角转旋转矩阵，ZYX顺序，单位为度"""
    rx, ry, rz = np.radians(euler)
    Rx = np.array([
        [1, 0, 0],
        [0, np.cos(rx), -np.sin(rx)],
        [0, np.sin(rx), np.cos(rx)]
    ])
    Ry = np.array([
        [np.cos(ry), 0, np.sin(ry)],
        [0, 1, 0],
        [-np.sin(ry), 0, np.cos(ry)]
    ])
    Rz = np.array([
        [np.cos(rz), -np.sin(rz), 0],
        [np.sin(rz), np.cos(rz), 0],
        [0, 0, 1]
    ])
    return Rz @ Ry @ Rx

def matrix_to_euler(R):
    """旋转矩阵转欧拉角（ZYX，度）"""
    sy = np.sqrt(R[0,0]**2 + R[1,0]**2)
    singular = sy < 1e-6
    if not singular:
        x = np.arctan2(R[2,1], R[2,2])
        y = np.arctan2(-R[2,0], sy)
        z = np.arctan2(R[1,0], R[0,0])
    else:
        x = np.arctan2(-R[1,2], R[1,1])
        y = np.arctan2(-R[2,0], sy)
        z = 0
    return np.degrees([x, y, z])

def residual_func(params, Otw_list, P3w_list, image_points_list, object_points, R_tc_list, T_tc_list):
    """计算残差函数，使用CSV文件中提供的靶标坐标系变换矩阵"""
    fx, fy, cx, cy = params[:4]
    euler = params[4:7]
    print("euler", euler)
    T_wc = params[7:10]
    R_wc = euler_to_matrix(euler)
    K = np.array([[fx, 0, cx],
                  [0, fy, cy],
                  [0,  0,  1]])
    
    res = []
    origin_errors = []  # 存储原点误差
    p3_errors = []     # 存储三号点误差
    reproj_errors = [] # 存储重投影误差
    
    for i, (Otw, P3w, img_pts, R_tc, T_tc) in enumerate(zip(Otw_list, P3w_list, image_points_list, R_tc_list, T_tc_list)):
        # 确保object_points是float32类型的numpy数组
        object_points = np.array(object_points, dtype=np.float32)
        img_pts = np.array(img_pts, dtype=np.float32)
        
        # 计算靶标点在相机坐标系下的位置
        P_c = R_wc @ object_points.T + T_wc[:, np.newaxis]
        
        # 投影到图像平面
        P_2d = K @ P_c
        P_2d = P_2d[:2] / P_2d[2]
        
        # 计算重投影误差
        reproj_error = np.linalg.norm(P_2d.T - img_pts, axis=1)
        reproj_errors.append(np.mean(reproj_error))
        
        # 计算靶标原点在相机坐标系下的位置
        Otc = R_wc @ Otw + T_wc
        
        # 计算靶标三号点在相机坐标系下的位置
        P3c = R_wc @ P3w + T_wc
        
        # 计算靶标三号点在靶标坐标系下的位置
        P3t = object_points[2]  # 三号点在靶标坐标系下的位置
        
        # 使用CSV文件中提供的变换矩阵计算三号点在相机坐标系下的位置
        P3c_estimated = R_tc @ P3t + T_tc
        
        # 计算误差
        origin_error = Otc - T_tc
        p3_error = P3c - P3c_estimated
        
        # 存储误差
        origin_errors.append(np.linalg.norm(origin_error))
        p3_errors.append(np.linalg.norm(p3_error))
        
        # 残差 = 重投影误差 + 几何一致性约束
        res.extend(reproj_error)  # 重投影误差
        res.extend(origin_error)  # 靶标原点几何一致性约束
        res.extend(p3_error)  # 靶标三号点几何一致性约束
    
    # 打印误差统计信息
    if len(res) > 0 and np.random.random() < 0.01:  # 随机打印，避免输出太多
        print("\n=== 残差分析 ===")
        print(f"重投影误差均值: {np.mean(reproj_errors):.2f} 像素")
        print(f"原点误差均值: {np.mean(origin_errors):.2f} mm")
        print(f"原点误差最大值: {np.max(origin_errors):.2f} mm")
        print(f"三号点误差均值: {np.mean(p3_errors):.2f} mm")
        print(f"三号点误差最大值: {np.max(p3_errors):.2f} mm")
        print(f"总残差范数: {np.linalg.norm(res):.2f}")
    
    return np.array(res)

def optimize_camera_params(initial_params, Otw_list, P3w_list, image_points_list, object_points, R_tc_list, T_tc_list):
    """优化相机参数"""
    # 定义参数边界
    bounds = (
        [0, 0, 0, 0, -np.inf, -np.inf, -np.inf, -np.inf, -np.inf, -np.inf],  # 下界
        [np.inf, np.inf, 1920, 1080, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf]  # 上界
    )
    
    # 使用最小二乘法优化
    result = least_squares(
        residual_func,
        initial_params,
        args=(Otw_list, P3w_list, image_points_list, object_points, R_tc_list, T_tc_list),
        bounds=bounds,
        method='trf',
        verbose=1
    )
    
    return result.x
